package datastructure.segmenttree;

/**
 * 处理区间问题
 */
public class SegmentTree<E> {

    private E[] tree;

    private E[] data;

    private Merger<E> merger;

    public SegmentTree(E[] arr, Merger<E> merger) {
        this.merger = merger;
        data = (E[]) new Object[arr.length];
        for (int i = 0; i < arr.length; i++) {
            data[i] = arr[i];
        }
        tree = (E[]) new Object[(data.length << 2)];
        buildSegmentTree(0, 0, data.length - 1);
    }

    // 在index位置创建表示区间[l...r]的线段树
    // 递归终止条件：只有1个元素，即左右孩子index相等，长度为1
    private void buildSegmentTree(int index, int l, int r) { // index:线段树结点，l数组左边，r数组右边
        if (l == r) {
            // 左右相等时，直接取值即可， 长度为1
            tree[index] = data[l];
            return;
        }
        int leftChild = leftChildIndex(index); // 线段树的左孩子
        int rightChild = rightChildIndex(index); // 线段树的右孩子

        int mid = l + (r - l) / 2; // 平分的中点位置
        buildSegmentTree(leftChild, l, mid);
        buildSegmentTree(rightChild, mid + 1, r);
        tree[index] = merger.merge(tree[leftChild], tree[rightChild]);
    }

    public E get(int index) {
        if (index < 0 || index >= data.length) {
            throw new IllegalArgumentException("");
        }
        return data[index];
    }

    private int leftChildIndex(int index) {
        return 2 * index + 1;
    }

    private int rightChildIndex(int index) {
        return 2 * index + 2;
    }
    // 线段树查询 例如：查询2~5的和

    /**
     * 根节点是0~7,要查询2~5
     * 分割点：l + (r-l)/2 = 3
     * 所以查询的是2~3和4~5
     * <p>
     * 根结点是0~3，要查询2~3
     * 分割点是1
     * 查询的是0~1和2~3
     * <p>
     * 右树，4~7
     */

    public E query(int queryL, int queryR) {
        if (queryL < 0 || queryR >= data.length) {
            throw new IllegalArgumentException("");
        }
        return query(0, 0, data.length - 1, queryL, queryR);
    }

    // 返回以treeIndex为根的线段树中[l.r]范围内，搜索区间[qL,qR]的值
    private E query(int treeIndex, int l, int r, int qL, int qR) {
        if (l == qL && r == qR) {
            return tree[treeIndex];
        }
        int mid = l + (r - l) / 2;
        int leftChild = leftChildIndex(treeIndex);
        int rightChild = rightChildIndex(treeIndex);
        if (qL >= mid + 1) {
            return query(rightChild, mid + 1, r, qL, qR);
        }
        if (qR <= mid) {
            return query(leftChild, l, mid, qL, qR);
        }
        E leftRes = query(leftChild, l, mid, qL, mid);
        E rightRes = query(rightChild, mid + 1, r, mid + 1, qR);
        return merger.merge(leftRes, rightRes);
    }

    /**
     * @param index
     * @param e
     */
    public void set(int index, E e) {
        {
            data[index] = e;
            set(0, 0, data.length - 1, index, e);
        }

    }

    private void set(int treeIndex, int l, int r, int index, E e) {
        if (l == r) {
            tree[index] = e;
            return;
        }
        int mid = l + (r - l) / 2;
        int leftChild = leftChildIndex(treeIndex);
        int rightChild = rightChildIndex(treeIndex);

        if (index >= mid + 1) {
            set(rightChild, mid + 1, r, index, e);
        } else {
            set(leftChild, l, mid, index, e);
        }
        tree[treeIndex] = merger.merge(tree[leftChild], tree[rightChild]);
    }
}
